# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

import torch
import torch_npu
import triton.language as tl
import triton
# from vllm.triton_utils import tl, triton


# 入口函数
'''
query: [num_tokens, num_query_heads, head_size]
key: [num_tokens, num_kv_heads, head_size]
value: [num_tokens, num_kv_heads, head_size]
'''
def blocksparse_flash_attn_varlen_fwd(
        q,
        k,
        v,  # (#tokens, n_heads, head_size)
        cu_seqlens_k,
        cu_seqlens_q,
        sm_scale,
        sparse_layout,
        *,
        block_size=64,
        q_block_size=None,
        max_seqlen=None):
    # split q to blocks

    assert isinstance(sparse_layout, (list, tuple))

    _, n_heads, head_size = q.shape
    # 获取批次数
    batch_size = cu_seqlens_k.size(0) - 1
    q_block_size = q_block_size or block_size

    assert q.dim() == k.dim() == v.dim() == 3
    # 必须是整数倍 这样才能均匀分配查询头
    assert q.size(1) % k.size(1) == 0
    assert q.size(2) == k.size(2)
    # TODO(linxihui): allow k, v to have different head_size
    assert k.shape == v.shape
    assert cu_seqlens_k.dim() == 1

    q_k_ratio = q.size(1) // k.size(1)

    if cu_seqlens_q is None:
        # 这里是判断是否是自解码阶段
        # 每个批次q_len=1，只生成一个新token
        if q.size(0) == batch_size:  # decoding only
            cu_seqlens_q = torch.arange(
                0,
                batch_size + 1,
                dtype=cu_seqlens_k.dtype,
                device=cu_seqlens_k.device,
            )
        # q和k的tokens的长度相同
        # 表示现在处于纯预填充阶段
        # 整个序列一次性计算
        # 完全一样就行
        elif q.size(0) == k.size(0):
            cu_seqlens_q = cu_seqlens_k
        else:
            raise ValueError("cu_seqlens_q must be specified\
                    if it mix of prefilling and decoding.")
    else:
        assert cu_seqlens_k.size(0) == cu_seqlens_q.size(0)

    # switch to use cpu to avoid too many kernel launches when iterated over
    # 获取q的间隔长度 以及 k的间隔长度
    q_lens = (cu_seqlens_q[1:] - cu_seqlens_q[:-1]).cpu()
    k_lens = (cu_seqlens_k[1:] - cu_seqlens_k[:-1]).cpu()

    # 进行判断，判断是否处于自解码阶段(Decoding)还是预填充阶段(Prefilling)
    assert torch.logical_or(q_lens == 1, k_lens == q_lens).all(), (
        "length of q should either be 1 (decoding) or same as k (prefilling).")

    if max_seqlen:
        assert k_lens.max() <= max_seqlen

    # q_lens按照q_block_size进行切分操作
    n_blocks = (q_lens + q_block_size - 1) // q_block_size

    # 一个长度为总块数的张量，每个元素表示该查询块所属的批次ID
    # n_blocks 按照q_lens进行了分块处理
    q_batch_ids = torch.tensor(
        # n是重复次数，因为q_lens是array类型的
        [i for i, n in enumerate(n_blocks) for _ in range(n)],
        dtype=cu_seqlens_q.dtype,
        device=cu_seqlens_q.device,
    )
    # 一个长度相同的张量，每个元素表示该查询块的起始存储ID
    # i * q_block_size
    q_start_sids = torch.tensor(
        [i * q_block_size for n in n_blocks for i in range(n)],
        dtype=cu_seqlens_q.dtype,
        device=cu_seqlens_q.device,
    )

    # 结果和q是一样的形状
    out = q.new_empty(q.shape)
    cu_seqlens_q = cu_seqlens_q.contiguous()
    cu_seqlens_k = cu_seqlens_k.contiguous()

    layout_crow_indices, layout_col_indices = sparse_layout
    # 对head_size获取下个2的阶乘
    block_d = triton.next_power_of_2(head_size)

    decoding_only = (q_lens == 1).all().item()
    # (total_q_blocks, n_heads, 1)
    # 每个线程块处理一个查询块的一个头 n_heads作为了grid的子单位
    grid = (len(q_start_sids), n_heads, 1)

    _fwd_kernel_batch_inference[grid](
        q,
        k,
        v,
        out,
        sm_scale,
        cu_seqlens_q[:-1],
        cu_seqlens_q[1:],
        cu_seqlens_k[:-1],
        cu_seqlens_k[1:],
        q_batch_ids,
        q_start_sids,
        0,
        *q.stride(),
        0,
        *k.stride(),
        0,
        *v.stride(),
        0,
        *out.stride(),
        layout_crow_indices,
        layout_col_indices,
        *layout_crow_indices.stride(),
        *layout_col_indices.stride(),
        q_k_ratio,
        HAS_BATCH_DIM=False,
        D_HEAD=head_size,
        BLOCK_M=q_block_size,
        BLOCK_N=block_size,
        BLOCK_D=block_d,
        BLOCK_M_LOADING=(16 if decoding_only else
                         q_block_size),  # smaller for decoding
        EVEN_D=block_d == head_size,
        num_warps=1 if decoding_only else 4,
        num_stages=3)

    return out

# 进行子块的qkv计算
@triton.jit
def _fwd_kernel_inner(
    acc,
    l_i,
    m_i,
    q,
    Q,
    k_block_col_idx,
    layout_col_ptr,
    layout_col_stride_h,
    layout_col_stride_m,
    k_ptrs,
    v_ptrs,
    off_h,
    offs_m,
    offs_n,
    offs_d,
    stride_kt,
    stride_vt,
    sm_scale,
    k_seqlen,
    past_len,
    LAST_K_BLOCK: tl.constexpr,
    BLOCK_M_LOADING: tl.constexpr,
    BLOCK_N: tl.constexpr,
    D_HEAD: tl.constexpr,
    EVEN_D: tl.constexpr,
    M_LT_N: tl.constexpr,
):
    # 获取当前处理k的块id
    k_block_id = tl.load(layout_col_ptr + off_h * layout_col_stride_h +
                         k_block_col_idx * layout_col_stride_m).to(tl.int32)
    # 起始地址
    start_n = k_block_id * BLOCK_N
    # 是否是尾块，尾块处理逻辑不一样
    if LAST_K_BLOCK:
        if EVEN_D:
            # 读取K的子块位置
            k = tl.load(
                k_ptrs + start_n * stride_kt,
                mask=offs_n[None, :] + start_n < k_seqlen,
                other=0.0,
            )
        else:
            k = tl.load(
                k_ptrs + start_n * stride_kt,
                mask=(offs_n[None, :] + start_n < k_seqlen) &
                (offs_d[:, None] < D_HEAD),
                other=0.0,
            )
    else:
        if EVEN_D:
            k = tl.load(k_ptrs + start_n * stride_kt)
        else:
            k = tl.load(k_ptrs + start_n * stride_kt,
                        mask=offs_d[:, None] < D_HEAD,
                        other=0.0)

    qk = tl.zeros([BLOCK_M_LOADING, BLOCK_N], dtype=tl.float32)
    qk += tl.dot(q, k)
    qk *= sm_scale

    # the following is needed only when LAST_K_BLOCK or BLOCK_M < BLOCK_N
    if LAST_K_BLOCK | M_LT_N:
        # 设置casual mask 确保上三角为-inf 设置计算掩码
        qk = tl.where(
            offs_m[:, None] + past_len >= (start_n + offs_n[None, :]),
            qk,
            float("-inf"),
        )

    # flash-attn2
    # 集成了flash-attention2的算法
    # 获取新的最大值
    m_ij = tl.maximum(m_i, tl.max(qk, 1))
    # 进行softmax操作的归一化处理
    p = tl.exp(qk - m_ij[:, None])
    # 进行子块加和操作 axis=1
    l_ij = tl.sum(p, 1)
    # 获取补偿系数
    alpha = tl.exp(m_i - m_ij)
    # p的原因 p进行exp操作
    acc = acc * alpha[:, None]
    # update m_i
    # 更新当前子块的最大值
    m_i = m_ij
    # 更新当前行的sum值
    l_i = l_i * alpha + l_ij

    p = p.to(Q.dtype.element_ty)
    # update acc
    # 尾块处理逻辑，省掉额外的逻辑处理开销
    if LAST_K_BLOCK:
        if EVEN_D:
            v = tl.load(
                v_ptrs + start_n * stride_vt,
                mask=offs_n[:, None] + start_n < k_seqlen,
                other=0.0,
            )
        else:
            v = tl.load(
                v_ptrs + start_n * stride_vt,
                mask=(offs_n[:, None] + start_n < k_seqlen) &
                (offs_d[None, :] < D_HEAD),
                other=0.0,
            )
    else:
        if EVEN_D:
            v = tl.load(v_ptrs + start_n * stride_vt)
        else:
            v = tl.load(v_ptrs + start_n * stride_vt,
                        mask=offs_d[None, :] < D_HEAD,
                        other=0.0)

    acc += tl.dot(p, v)

    return acc, l_i, m_i


# 这里来进行判断，从而来设置参数
# 判断M轴长和N轴长之间的关系
@triton.heuristics({
    "M_LT_N":
    lambda kwargs: kwargs["BLOCK_M"] < kwargs["BLOCK_N"],
})
@triton.jit
def _fwd_kernel_batch_inference(
    Q, # [total_tokens, n_query_heads, head_size]
    K, # [total_tokens, n_kv_heads, head_size]
    V, # [total_tokens, n_kv_heads, head_size]
    Out, # [total_tokens, n_query_heads, head_size]
    sm_scale, # 缩放参数
    q_batch_starts, # [batch_size]
    q_batch_ends, # [batch_size]
    k_batch_starts, # [batch_size]
    k_batch_ends, # [batch_size]
    q_batch_ids, # [total_q_blocks]
    q_start_sids, # [total_q_blocks]
    # 这些通过*操作符来进行解码操作
    stride_qb, # query batch_stride
    stride_qt, # query tokens_stride
    stride_qh, # query heads_stride
    stride_qd, # query D_stride
    stride_kb, # key batch_stride
    stride_kt, # key tokens_stride
    stride_kh, # key heads_stride
    stride_kd, # key D_stride
    stride_vb, # value batch_stride
    stride_vt, # value tokens_stride
    stride_vh, # value heads_stride
    stride_vd, # value D_stride
    stride_ob, # out batch_stride
    stride_ot, # out tokens_stride
    stride_oh, # out heads_stride
    stride_od, # out D_stride
    layout_crow_ptr, # [n_heads, num_q_blocks + 1]
    layout_col_ptr, # [n_heads, total_valid_blocks]
    layout_crow_stride_h,
    layout_crow_stride_m,
    layout_col_stride_h,
    layout_col_stride_m,
    q_k_ratio, # 查询/键头比率
    HAS_BATCH_DIM: tl.constexpr,
    D_HEAD: tl.constexpr,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
    BLOCK_D: tl.constexpr,
    BLOCK_M_LOADING: tl.constexpr,
    # EVEN_D 可以作为对齐操作，mask特殊判断
    EVEN_D: tl.constexpr,
    M_LT_N: tl.constexpr, # 查询块 < 键块
):
    """
    NOTATION:
    pid: position id
    sid: storage id
    sbid: storage block id
    pbid: position block id
    offs_m, offs_n: storage offsets of m-dim(q, row) and n-dim(k, col)

    TODO(linxihui):
    Optimize grouped-attn
    """
    # 按照total_tokens来进行分配线程
    # off_zm 是对应查询块索引
    off_zm = tl.program_id(0)
    # off_h 是对应查询头索引
    off_h = tl.program_id(1)

    # 针对head进行偏移计算
    off_h_for_kv = off_h // q_k_ratio

    if HAS_BATCH_DIM:
        # 获取批次信息
        off_z = tl.program_id(2)
        Q += off_z * stride_qb
        K += off_z * stride_kb
        V += off_z * stride_vb
        Out += off_z * stride_ob
        start_m = off_zm
        q_start_sid = start_m * BLOCK_M  # always 0 for decoding
    else:
        # 标识当前查询块属于哪个批次
        off_z = tl.load(q_batch_ids + off_zm).to(tl.int32)  # [0, 0, 0, 1]
        # 获取每个查询块起始token位置  绝对位置
        q_start_sid = tl.load(q_start_sids + off_zm)
        # 获取当前查询块的分块位置 相对位置
        start_m = q_start_sid // BLOCK_M  # q_sbid

    # 获取当前处理块的偏移量
    # 存在两个BLOCK_M 以及 BLOCK_M_LOADING
    offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M_LOADING)
    # 键序列块内方向
    offs_n = tl.arange(0, BLOCK_N)
    # 头维度信息
    # BLOCK_D 是 head_size 的下一个 2 的幂（next_power_of_2(head_size)），用于对齐内存访问
    offs_d = tl.arange(0, BLOCK_D)

    # 获取qk的起始位置以及末尾位置，获取长度
    q_cu_start = tl.load(q_batch_starts + off_z).to(tl.int32)
    q_seqlen = tl.load(q_batch_ends + off_z).to(tl.int32) - q_cu_start
    k_cu_start = tl.load(k_batch_starts + off_z).to(tl.int32)
    k_seqlen = tl.load(k_batch_ends + off_z).to(tl.int32) - k_cu_start
    # 过去已生成的token长度
    # 在decoding阶段，q_seqlen = 1 k_seqlen = 当前总序列长度
    past_len = k_seqlen - q_seqlen

    # 当前线程块处理的指针位置
    Q += q_cu_start * stride_qt + off_h * stride_qh
    # off_h_for_kv 存在共享头信息
    K += k_cu_start * stride_kt + off_h_for_kv * stride_kh
    V += k_cu_start * stride_vt + off_h_for_kv * stride_vh
    # Out最后的形状和Q一样
    Out += q_cu_start * stride_ot + off_h * stride_oh

    q_pbid = (past_len + q_start_sid) // BLOCK_M

    if EVEN_D:
        # 子块位置
        q = tl.load(
            Q + offs_m[:, None] * stride_qt + offs_d[None, :] * stride_qd,
            mask=offs_m[:, None] < q_seqlen,
            other=0.0,
        )
    else:
        q = tl.load(
            Q + offs_m[:, None] * stride_qt + offs_d[None, :] * stride_qd,
            # offs_d[None, :] 来特殊判断是否越界
            mask=(offs_m[:, None] < q_seqlen) & (offs_d[None, :] < D_HEAD),
            other=0.0,
        )

    sparse_crow_ptr = (layout_crow_ptr + off_h * layout_crow_stride_h +
                       q_pbid * layout_crow_stride_m)

    # TODO(linxihui): load at once, with any Triton version
    # that supports `tl.split`, e.g., Triton 3.0
    k_block_start = tl.load(sparse_crow_ptr).to(tl.int32)
    k_block_end = tl.load(sparse_crow_ptr + 1).to(tl.int32)

    # 根据flash attention2 设置比较值
    # 记录当前行的最大值 max logit
    m_i = tl.zeros([BLOCK_M_LOADING], dtype=tl.float32) - float("inf")
    # 存储当前行 sum(exp(logit - m_i)) 进行逐步更新
    l_i = tl.zeros([BLOCK_M_LOADING], dtype=tl.float32)
    # qkv的乘法的累计值
    acc = tl.zeros([BLOCK_M_LOADING, BLOCK_D], dtype=tl.float32)

    # K和V的形状相同 所以是一样加载方式
    k_ptrs = K + offs_n[None, :] * stride_kt + offs_d[:, None] * stride_kd
    v_ptrs = V + offs_n[:, None] * stride_vt + offs_d[None, :] * stride_vd

    # 缩放参数
    sm_scale *= (
        1.44269504  # 1/log2 as we use base2 for exponential and logarithm
    )

    # 进行循环处理
    for k_block_col_idx in range(k_block_start, k_block_end - 1):
        acc, l_i, m_i = _fwd_kernel_inner(
            acc,
            l_i,
            m_i,
            q,
            Q,
            k_block_col_idx,
            layout_col_ptr,
            layout_col_stride_h,
            layout_col_stride_m,
            k_ptrs,
            v_ptrs,
            off_h,
            offs_m,
            offs_n,
            offs_d,
            stride_kt,
            stride_vt,
            sm_scale,
            k_seqlen,
            past_len,
            False,
            BLOCK_M_LOADING,
            BLOCK_N,
            D_HEAD,
            EVEN_D,
            M_LT_N,
        )

    # 尾块处理
    acc, l_i, m_i = _fwd_kernel_inner(
        acc,
        l_i,
        m_i,
        q,
        Q,
        k_block_end - 1,
        layout_col_ptr,
        layout_col_stride_h,
        layout_col_stride_m,
        k_ptrs,
        v_ptrs,
        off_h,
        offs_m,
        offs_n,
        offs_d,
        stride_kt,
        stride_vt,
        sm_scale,
        k_seqlen,
        past_len,
        True,
        BLOCK_M_LOADING,
        BLOCK_N,
        D_HEAD,
        EVEN_D,
        M_LT_N,
    )

    # flash-attn 2
    m_i += tl.math.log2(l_i)
    acc = acc / l_i[:, None]

    # write output
    if EVEN_D:
        tl.store(
            Out + offs_m[:, None] * stride_ot + offs_d[None, :] * stride_od,
            acc,
            mask=offs_m[:, None] < q_seqlen,
        )
    else:
        tl.store(
            Out + offs_m[:, None] * stride_ot + offs_d[None, :] * stride_od,
            acc,
            mask=(offs_m[:, None] < q_seqlen) & (offs_d[None, :] < D_HEAD),
        )